import onnx.hub
import pyrealsense2 as rs
import numpy as np
import cv2
import torch
import threading
import pathlib
temp = pathlib.PosixPath
pathlib.PosixPath = pathlib.WindowsPath


model = torch.hub.load('../yolov5','custom', path='../yolov5/best.pt', source='local', device='cpu')
model.conf = 0.5
background = np.zeros((960, 1280, 3), dtype=np.uint8)


def dectshow(org_img, boxs):
    img = org_img.copy()
    for box in boxs:
        cv2.rectangle(img, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 255, 0), 2)
        cv2.putText(img, str(box[-1]), (int(box[0]), int(box[1])), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
    return img


def process_device(device, window_name, index):
    pipeline = rs.pipeline()
    config = rs.config()
    config.enable_device(device.get_info(rs.camera_info.serial_number))
    config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 60)
    pipeline.start(config)
    try:
        while True:
            frames = pipeline.wait_for_frames()
            color_frame = frames.get_color_frame()
            color_image = np.asanyarray(color_frame.get_data())
            # color_image_resized = cv2.resize(color_image, (640, 640))
            results = model(color_image)
            boxs= results.pandas().xyxy[0].values
            image_with_boxes = dectshow(color_image, boxs)

            x = (index % 2) * 640
            y = (index // 2) * 480
            background[y:y+480, x:x+640] = image_with_boxes

            cv2.imshow("Traffic System", background)

            key = cv2.waitKey(1)
            if key & 0xFF == ord('q') or key == 27:
                break
    finally:
        pipeline.stop()


if __name__ == "__main__":
    ctx = rs.context()
    devices = ctx.query_devices()
    threads = []

    for i, device in enumerate(devices):
        window_name = f"Device {i + 1}: Serial Number - {device.get_info(rs.camera_info.serial_number)}"
        thread = threading.Thread(target=process_device, args=(device, window_name, i))
        threads.append(thread)

    for thread in threads:
        thread.start()

    for thread in threads:
        thread.join()

    cv2.destroyAllWindows()
